import os
from tqdm import tqdm

class DatasetDeleter:
    def __init__(self, dataset_root, prefix, subsets_to_scan):
        self.dataset_root = dataset_root
        self.prefix = prefix
        self.subsets_to_scan = subsets_to_scan
    def _delete_files_in_subset(self, subset_path):
        deleted_count = 0
        for root, dirs, files in os.walk(subset_path):
            # 使用 tqdm 包装文件列表
            for file in tqdm(files, desc=f"处理 {os.path.basename(root)} 目录文件", unit="file"):
                if file.startswith(self.prefix):
                    file_path = os.path.join(root, file)
                    try:
                        os.remove(file_path)
                        deleted_count += 1
                    except Exception as e:
                        print(f"删除文件 {file_path} 时出错: {e}")
        print(f"在 {subset_path} 中总共删除了 {deleted_count} 个文件。")
    def run(self):
        for subset in self.subsets_to_scan:
            subset_path = os.path.join(self.dataset_root, subset)
            if os.path.exists(subset_path):
                self._delete_files_in_subset(subset_path)


if __name__ == "__main__":
    deleter = DatasetDeleter(
        dataset_root='datasets_balanced',
        prefix='630630',
        subsets_to_scan=['train', 'valid', 'test']
    )
    deleter.run()